import torch

# 定义调度器类型常量
SCHEDULER_TYPE_BATCH = 'batch'  # 每个batch后调用step
SCHEDULER_TYPE_EPOCH = 'epoch'  # 每个epoch后调用step
SCHEDULER_TYPE_VAL = 'val'      # 验证后调用step并传入指标

# 调度器类型映射表
SCHEDULER_TYPE_MAP = {
    # 每个batch后调用step的调度器
    'OneCycleLR': SCHEDULER_TYPE_BATCH,
    'CyclicLR': SCHEDULER_TYPE_BATCH,
    'CosineAnnealingWarmRestarts': SCHEDULER_TYPE_BATCH,
    'ExponentialLR': SCHEDULER_TYPE_BATCH,  # 也可在epoch后调用
    'StepLR': SCHEDULER_TYPE_BATCH,          # 也可在epoch后调用
    
    # 每个epoch后调用step的调度器
    'MultiStepLR': SCHEDULER_TYPE_EPOCH,
    'LinearLR': SCHEDULER_TYPE_EPOCH,
    'CosineAnnealingLR': SCHEDULER_TYPE_EPOCH,
    'PolynomialLR': SCHEDULER_TYPE_EPOCH,
    
    # 需要验证指标的调度器
    'ReduceLROnPlateau': SCHEDULER_TYPE_VAL,
    'RAdam': SCHEDULER_TYPE_VAL,  # 示例，实际可能不同
    'CosineAnnealingLRWithWarmup': SCHEDULER_TYPE_VAL,  # 自定义调度器示例
}


class SchedulerManager:
    """
    学习率调度器管理器，统一管理不同类型调度器的step()调用位置
    """
    def __init__(self, scheduler, scheduler_type=None):
        """
        初始化调度器管理器
        
        参数:
            scheduler: 学习率调度器实例
            scheduler_type: 调度器类型，可以是'batch'、'epoch'或'val'。
                          如果为None，则自动根据调度器类名确定
        """
        self.scheduler = scheduler
        self.scheduler_class_name = scheduler.__class__.__name__
        
        # 自动确定调度器类型
        if scheduler_type is None:
            self.scheduler_type = SCHEDULER_TYPE_MAP.get(
                self.scheduler_class_name, SCHEDULER_TYPE_EPOCH)
        else:
            self.scheduler_type = scheduler_type
        
        print(f"调度器管理器初始化: {self.scheduler_class_name} 类型: {self.scheduler_type}")
    
    def step(self, metric=None):
        """
        根据调度器类型执行相应的step()调用
        
        参数:
            metric: 验证指标，仅对需要验证指标的调度器有效
        
        返回:
            bool: 是否成功执行step()
        """
        try:
            if self.scheduler_type == SCHEDULER_TYPE_VAL:
                # 需要验证指标的调度器
                if metric is not None:
                    self.scheduler.step(metric)
                    return True
                else:
                    print(f"警告: 调度器 {self.scheduler_class_name} 需要验证指标，但未提供")
                    return False
            else:
                # 不需要验证指标的调度器
                self.scheduler.step()
                return True
        except Exception as e:
            print(f"执行调度器step()时出错: {str(e)}")
            return False
    
    def get_lr(self):
        """
        获取当前学习率
        
        返回:
            list: 当前学习率列表
        """
        if hasattr(self.scheduler, 'get_last_lr'):
            return self.scheduler.get_last_lr()
        elif hasattr(self.scheduler.optimizer, 'param_groups'):
            return [group['lr'] for group in self.scheduler.optimizer.param_groups]
        else:
            return [self.scheduler.optimizer.lr]
    
    def adjust_state(self, epoch):
        """调整调度器状态到指定的epoch，避免直接调用step()导致的警告
        
        参数:
        - epoch: 要调整到的epoch数
        """
        if self.scheduler:
            # 对于不同类型的调度器，使用不同的方法来设置状态
            if hasattr(self.scheduler, '_step_count'):
                # 对于大多数调度器，直接设置_step_count
                self.scheduler._step_count = epoch
            elif hasattr(self.scheduler, 'last_epoch'):
                # 对于某些调度器（如ReduceLROnPlateau），设置last_epoch
                self.scheduler.last_epoch = epoch - 1  # 因为大多数调度器会在step()时增加last_epoch
    
    def is_batch_level(self):
        """
        检查是否是batch级别的调度器
        
        返回:
            bool: 是否是batch级别
        """
        return self.scheduler_type == SCHEDULER_TYPE_BATCH
    
    def is_val_level(self):
        """
        检查是否是需要验证指标的调度器
        
        返回:
            bool: 是否需要验证指标
        """
        return self.scheduler_type == SCHEDULER_TYPE_VAL
    
    def needs_batch_update(self):
        """
        检查是否需要在每个batch后更新学习率
        
        返回:
            bool: 是否需要batch级别的更新
        """
        return self.is_batch_level()
    
    def needs_epoch_update(self):
        """
        检查是否需要在每个epoch后更新学习率
        
        返回:
            bool: 是否需要epoch级别的更新
        """
        return self.scheduler_type == SCHEDULER_TYPE_EPOCH
    
    def needs_val_update(self):
        """
        检查是否需要在验证后更新学习率
        
        返回:
            bool: 是否需要验证后的更新
        """
        return self.is_val_level()
    
    def step_batch(self):
        """
        在每个batch后执行调度器的step操作
        这是为了兼容trainer.py中的调用
        
        返回:
            bool: 是否成功执行step()
        """
        return self.step()
    
    def step_epoch(self):
        """
        在每个epoch后执行调度器的step操作
        这是为了兼容trainer.py中的调用
        
        返回:
            bool: 是否成功执行step()
        """
        return self.step()
    
    def step_val(self, metric):
        """
        在验证后执行调度器的step操作，并传入验证指标
        这是为了兼容trainer.py中的调用
        
        参数:
            metric: 验证指标值
            
        返回:
            bool: 是否成功执行step()
        """
        return self.step(metric)


def create_scheduler(optimizer, config):
    """
    创建学习率调度器和对应的管理器
    确保正确支持sweep_config.yaml中配置的所有调度器类型
    
    参数:
        optimizer: 优化器实例
        config: 配置字典
    
    返回:
        tuple: (scheduler_manager, scheduler)
    """
    if 'lr_scheduler' not in config:
        print("配置文件中缺少 'lr_scheduler' 字段，不使用学习率调度器")
        return None, None
    
    scheduler_config = config['lr_scheduler']
    scheduler_type = scheduler_config['type']
    scheduler_args = scheduler_config.get('args', {})
    
    # 获取训练参数
    train_epochs = config.get('trainer', {}).get('epochs', 100)
    batch_size = config.get('data_loader', {}).get('batch_size', 128)
    
    # 针对不同调度器的特殊参数处理
    # 1. CosineAnnealingLR
    if scheduler_type == 'CosineAnnealingLR':
        if 'T_max' not in scheduler_args:
            scheduler_args['T_max'] = train_epochs
            print(f"自动设置CosineAnnealingLR的T_max为训练总轮数: {scheduler_args['T_max']}")
    
    # 2. CosineAnnealingWarmRestarts
    elif scheduler_type == 'CosineAnnealingWarmRestarts':
        if 'T_0' not in scheduler_args:
            scheduler_args['T_0'] = max(10, train_epochs // 5)  # 默认为训练轮数的1/5
            print(f"自动设置CosineAnnealingWarmRestarts的T_0: {scheduler_args['T_0']}")
    
    # 3. OneCycleLR - 需要特殊处理，通常需要steps_per_epoch参数
    elif scheduler_type == 'OneCycleLR':
        # 尝试获取数据加载器信息来估算steps_per_epoch
        if 'steps_per_epoch' not in scheduler_args and 'data_loader' in config:
            # 如果没有提供，设置为一个合理的默认值
            # 假设每个epoch有100个steps
            scheduler_args['steps_per_epoch'] = 100
            print(f"自动设置OneCycleLR的steps_per_epoch为: {scheduler_args['steps_per_epoch']}")
        
        # 设置epochs参数
        if 'epochs' not in scheduler_args:
            scheduler_args['epochs'] = train_epochs
            print(f"自动设置OneCycleLR的epochs为: {scheduler_args['epochs']}")
        
        # 如果没有提供max_lr，基于优化器当前lr设置
        if 'max_lr' not in scheduler_args:
            current_lr = optimizer.param_groups[0]['lr']
            scheduler_args['max_lr'] = current_lr * 10  # 默认为当前lr的10倍
            print(f"自动设置OneCycleLR的max_lr为: {scheduler_args['max_lr']}")
    
    # 4. ReduceLROnPlateau
    elif scheduler_type == 'ReduceLROnPlateau':
        # 确保mode参数存在
        if 'mode' not in scheduler_args:
            scheduler_args['mode'] = 'min'  # 默认监控最小值
        
        # 设置默认patience
        if 'patience' not in scheduler_args:
            scheduler_args['patience'] = 5
    
    # 5. MultiStepLR
    elif scheduler_type == 'MultiStepLR':
        # 如果没有提供milestones，设置一些默认值
        if 'milestones' not in scheduler_args:
            scheduler_args['milestones'] = [int(train_epochs * 0.5), int(train_epochs * 0.75)]
            print(f"自动设置MultiStepLR的milestones为: {scheduler_args['milestones']}")
        
        # 确保gamma存在
        if 'gamma' not in scheduler_args:
            scheduler_args['gamma'] = 0.1
    
    # 6. LinearLR
    elif scheduler_type == 'LinearLR':
        # 确保start_factor和end_factor存在
        if 'start_factor' not in scheduler_args:
            scheduler_args['start_factor'] = 1.0
        if 'end_factor' not in scheduler_args:
            scheduler_args['end_factor'] = 0.01
    
    # 7. ExponentialLR
    elif scheduler_type == 'ExponentialLR':
        # 确保gamma存在
        if 'gamma' not in scheduler_args:
            scheduler_args['gamma'] = 0.95
    
    # 8. StepLR
    elif scheduler_type == 'StepLR':
        # 确保step_size和gamma存在
        if 'step_size' not in scheduler_args:
            scheduler_args['step_size'] = max(10, train_epochs // 5)
        if 'gamma' not in scheduler_args:
            scheduler_args['gamma'] = 0.1
    
    # 参数类型转换：将字符串形式的数值转换为相应的数值类型
    numeric_params = ['eta_min', 'T_max', 'T_0', 'T_mult', 'base_lr', 'max_lr', 'step_size', 
                     'gamma', 'patience', 'factor', 'min_lr', 'threshold', 'steps_per_epoch',
                     'epochs', 'start_factor', 'end_factor', 'pct_start']
    
    for param_name in scheduler_args:
        if param_name in numeric_params and isinstance(scheduler_args[param_name], str):
            try:
                scheduler_args[param_name] = float(scheduler_args[param_name])
            except ValueError:
                # 如果转换失败，检查是否是列表
                if scheduler_args[param_name].strip().startswith('[') and scheduler_args[param_name].strip().endswith(']'):
                    try:
                        # 尝试解析为列表
                        import ast
                        scheduler_args[param_name] = ast.literal_eval(scheduler_args[param_name])
                    except:
                        pass
    
    # 确保milestones是列表类型
    if scheduler_type in ['MultiStepLR'] and 'milestones' in scheduler_args:
        if isinstance(scheduler_args['milestones'], (int, float)):
            scheduler_args['milestones'] = [int(scheduler_args['milestones'])]
        elif isinstance(scheduler_args['milestones'], str):
            try:
                import ast
                scheduler_args['milestones'] = ast.literal_eval(scheduler_args['milestones'])
            except:
                print(f"警告: 无法解析milestones参数: {scheduler_args['milestones']}")
                scheduler_args['milestones'] = [int(train_epochs * 0.5), int(train_epochs * 0.75)]
    
    try:
        # 动态创建调度器
        scheduler_module = torch.optim.lr_scheduler
        
        # 检查调度器是否存在
        if not hasattr(scheduler_module, scheduler_type):
            raise ValueError(f"PyTorch中不存在学习率调度器: {scheduler_type}")
        
        scheduler_class = getattr(scheduler_module, scheduler_type)
        
        # 获取调度器构造函数参数
        import inspect
        sig = inspect.signature(scheduler_class.__init__)
        valid_params = list(sig.parameters.keys())
        
        # 过滤掉不支持的参数
        filtered_args = {}
        for key, value in scheduler_args.items():
            if key in valid_params or key == 'params':  # 'params'是优化器参数，不需要包含在args中
                filtered_args[key] = value
        
        # 创建调度器实例
        scheduler = scheduler_class(optimizer, **filtered_args)
        
        # 创建调度器管理器
        scheduler_manager = SchedulerManager(scheduler)
        
        print(f"成功创建学习率调度器: {scheduler_type}")
        print(f"使用的调度器参数: {filtered_args}")
        print(f"调度器类型: {scheduler_manager.scheduler_type}")
        
        return scheduler_manager, scheduler
    except Exception as e:
        raise ValueError(f"创建学习率调度器 {scheduler_type} 失败: {str(e)}")


def get_scheduler_info(scheduler_manager):
    """
    获取调度器信息
    
    参数:
        scheduler_manager: 调度器管理器实例
    
    返回:
        dict: 调度器信息字典
    """
    if scheduler_manager is None:
        return None
    
    return {
        'class_name': scheduler_manager.scheduler_class_name,
        'type': scheduler_manager.scheduler_type,
        'current_lr': scheduler_manager.get_lr()[0],
        'is_batch_level': scheduler_manager.is_batch_level(),
        'is_val_level': scheduler_manager.is_val_level()
    }